// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include "_flat_hash_map.h"

#include <tbir/runtime/object.h>
#include <tbir/runtime/runtime_value.h>

namespace tbir::runtime {

    /*! \brief Dict node content in Dict */
    class DictNode : public Object {
    public:
        // data holder
        using _key_type = typename std::remove_cv<typename std::remove_reference<RTValue>::type>::type;
        using _value_type = typename std::remove_cv<typename std::remove_reference<RTValue>::type>::type;
        // using container_type = std::unordered_map<_key_type, _value_type, RTValueHash, RTValueEqual>;
        using container_type = ska::flat_hash_map<_key_type, _value_type>;

    public:
        // types
        using reference = typename container_type::reference;
        using const_reference = typename container_type::const_reference;
        using iterator = typename container_type::iterator;
        using const_iterator = typename container_type::const_iterator;
        using size_type = typename container_type::size_type;
        using difference_type = typename container_type::difference_type;
        using value_type = typename container_type::value_type;
        using allocator_type = typename container_type::allocator_type;
        using pointer = typename container_type::pointer;
        using const_pointer = typename container_type::const_pointer;
        using key_type = typename container_type::key_type;
        using mapped_type = typename container_type::mapped_type;

    public:
        DictNode() : data_container() {
        }

        DictNode(std::initializer_list<value_type> init) : data_container(init) {
        }

        template<class B, class E>
        DictNode(B begin, E end) : data_container(begin, end) {
        }

        static constexpr const uint32_t _type_index = TypeIndex::kRuntimeDict;
        static constexpr const char *_type_key = "runtime.Dict";
        TBIR_DECLARE_FINAL_OBJECT_INFO(DictNode, Object);

    public:
        // iterators
        TURBO_FORCE_INLINE iterator begin() {
            return data_container.begin();
        }

        TURBO_FORCE_INLINE const_iterator begin() const {
            return data_container.begin();
        }

        TURBO_FORCE_INLINE iterator end() {
            return data_container.end();
        }

        TURBO_FORCE_INLINE const_iterator end() const {
            return data_container.end();
        }

    public:
        template<class KEY_T>
        TURBO_FORCE_INLINE const_iterator find(const KEY_T &key) const {
            return data_container.find(key);
        }

        template<class KEY_T>
        TURBO_FORCE_INLINE iterator find(const KEY_T &key) {
            return data_container.find(key);
        }

        template<typename It>
        TURBO_FORCE_INLINE void insert(It begin, It end) {
            return data_container.insert(begin, end);
        }

        TURBO_FORCE_INLINE void insert(std::initializer_list<value_type> il) {
            return data_container.insert(il);
        }

        TURBO_FORCE_INLINE mapped_type &operator[](key_type key) {
            return data_container[std::move(key)];
        }

        template<typename Key, typename... Args>
        TURBO_FORCE_INLINE std::pair<iterator, bool> emplace(Key &&key, Args &&... args) {
            return data_container.emplace(std::forward<Key>(key), std::forward<Args>(args)...);
        }

        TURBO_FORCE_INLINE void clear() {
            data_container.clear();
        }

        TURBO_FORCE_INLINE void reserve(int64_t new_size) {
            if (new_size > 0) {
                data_container.reserve(new_size);
            }
        }

    public:
        // const methods in std::unordered_map
        TURBO_FORCE_INLINE size_t size() const {
            return data_container.size();
        }

        TURBO_FORCE_INLINE int64_t bucket_count() const {
            return data_container.bucket_count();
        }

        TURBO_FORCE_INLINE bool empty() const {
            return data_container.empty();
        }

        template<class KEY_U>
        TURBO_FORCE_INLINE bool contains(KEY_U const &key) const {
            return data_container.find(key) != data_container.end();
        }

    private:
        container_type data_container;

        // Reference class
        friend class Dict;

        friend class DictNodeTrait;
    };

}  // namespace tbir::runtime
